"""
write by qianqianjun
命令行程序
设置超参数，运行参数，运行环境等
"""

import argparse
import os
import random
import numpy as np
import torch
import torch.backends.cudnn as cudnn

class Parameter(argparse.ArgumentParser):
    def __init__(self):
        """
        设置训练过程的超参数
        """
        super(Parameter, self).__init__()
        arg = argparse.ArgumentParser()
        self.add_argument("--workers", type=int, help="数据加载线程的个数，默认4 线程", default=4)
        self.add_argument("--batchSize", type=int, help="输入批数据的大小，默认1", default=1)
        self.add_argument("--lr", type=float, default=0.0002, help="学习率，默认大小是 0.0002")
        self.add_argument("--dirCheckpoints", default="./checkpoints", help="存放训练好的模型的目录", type=str)
        self.add_argument("--dirImageOutput", default="./imgout", help="图片输出目录", type=str)
        self.add_argument("--dirTestOutput", default="./testout", help="测试 结果或者图像 的目录", type=str)
        self.add_argument("--modelPath", type=str, default="", help="存放训练好的模型的目录")
        self.add_argument("--useCuda", default=True, help="是否使用 GPU 加速，默认不使用", type=bool)
        self.add_argument("--randomSeed", help="用于实验可重复性验证的随机种子", default=None, type=int)
        self.add_argument("--useDense", default=False, help="是否使用 dense net 架构", type=bool)
        self.add_argument("--beta1", default=0.5, type=float, help="Adam 优化器需要使用的参数")
        self.add_argument("--epochs", default=100, type=int, help="对整个数据集训练多少次")

        self.readCmd(self.parse_args(args=[]))
    def readCmd(self, args):
        """
        设置参数可以方便外部访问
        :param args: 命令行传过来的参数列表
        :return: None
        """
        self.workers=args.workers
        self.batchSize=args.batchSize
        self.learning_rate=args.lr
        self.dirCheckpoints=args.dirCheckpoints
        self.dirImageOutput=args.dirImageOutput
        self.dirTestOutput=args.dirTestOutput
        self.useCuda=args.useCuda
        self.randomSeed=args.randomSeed
        self.useDense=args.useDense
        self.modelPath=args.modelPath
        self.beta1=args.beta1
        self.epochs=args.epochs

    def setImageInfo(self,imgSize=64,ngf=32,ndf=32,channel=3):
        """
        设置图片信息相关的参数
        :param imgSize:  图片的大小（这里为边长，图片选用方形的）
        :param ngf:  起始卷积核个数（用于创建网络层结构）
        :param ndf:  结束卷积核个数（用于创建网络层结构）
        :param channel:  图片的通道数目
        :return: None
        """
        self.imgSize=imgSize
        self.ngf=ngf
        self.ndf=ndf
        self.channel=channel

    def setLantentInfo(self,idim=16,wdim=128,zdim=128):
        """
        设置潜在空间的相关信息
        :param idim:  纹理空间的维度
        :param wdim:  变形场空间的维度
        :param zdim:  通用纹理空间的维度
        :return:  None
        """
        self.idim=idim
        self.wdim=wdim
        self.zdim=zdim

    def setPretrainedModel(self,path="/home/qianqianjun/CODE/DataSets/vgg19.pth"):
        self.pretrainedModel=path

parameter=Parameter()
parameter.setImageInfo(imgSize=224)
parameter.setLantentInfo(idim=32,wdim=256,zdim=256)
parameter.setPretrainedModel()


# 创建运行环境
# 创建实验用到的目录环境
os.makedirs(parameter.dirCheckpoints,exist_ok=True)
os.makedirs(parameter.dirImageOutput,exist_ok=True)
os.makedirs(parameter.dirTestOutput,exist_ok=True)

# 实验可重复性设置（设置随机种子）
if parameter.randomSeed is None:
    parameter.randomSeed=random.randint(1,100000000)
print("本次实验使用的随机种子",parameter.randomSeed)

random.seed(parameter.randomSeed)
np.random.seed(parameter.randomSeed)
torch.manual_seed(parameter.randomSeed)

if parameter.useCuda:
    torch.cuda.manual_seed_all(parameter.randomSeed)

# 加速设置
cudnn.benchmark=True
# 强制使用 GPU 加速
if torch.cuda.is_available() and not parameter.useCuda:
    print("当前GPU 设备已就绪，请使用GPU加速: --useCuda=True")
    exit(0)